import numpy as np
import logging


class Fusion:
    def __init__(self, fov=(-45, 45), pad_val=0):
        """
        初始化融合类
        :param fov: 视场角度范围，默认为 (-45, 45) 度
        :pad_val: 针对雷达判定为障碍物后分割图填充数值
        """
        self.fov = fov
        self.pad_val = pad_val
        self.__log_init()

    def __log_init(self):
        self.logger = logging.getLogger("Fusion")

    def mask_obstacle_angles(self, res, obstacle_angles):
        """
        根据障碍物的角度范围, 将图像中对应区域的像素值设置为 0
        :param res: 分割结果，SegDataSample
        :param obstacle_angles: 障碍物的角度范围列表，每个元素是一个元组 (angle_start, angle_end)
        :return: 处理后的图像, numpy 数组
        """
        seg_image = np.array(res.pred_sem_seg.data.squeeze().cpu())
        seg_image = seg_image.astype(np.uint8)

        height, width = seg_image.shape
        fov_min, fov_max = self.fov

        for angle_start, angle_end in obstacle_angles:
            # 将角度范围限制在图像的水平视角范围内
            angle_start = max(angle_start, fov_min)
            angle_end = min(angle_end, fov_max)

            # 将角度转换为对应的列坐标（注意：左边是 +45 度，右边是 -45 度，线性映射）
            # 因此，angle 要从 (+45 ~ -45) 映射到 (0 ~ width-1)
            col_start = int((fov_max - angle_start) / (fov_max - fov_min) * (width - 1))
            col_end = int((fov_max - angle_end) / (fov_max - fov_min) * (width - 1))

            # 确保索引在图像宽度范围内，且从小到大
            col_min = min(max(0, col_start), width - 1)
            col_max = min(max(0, col_end), width - 1)

            if col_min > col_max:
                col_min, col_max = col_max, col_min

            # 打印调试信息
            # self.logger.verbose(
            #     f"columns from {col_min} to {col_max} for angles {angle_start} to {angle_end}"
            # )

            # 设置这几列的所有像素为 0
            seg_image[:, col_min : col_max + 1] = self.pad_val

        return seg_image
